#!/usr/bin/python
#-------------------------#
# pymd5hashreverse.py     #
#-------------------------#
import argparse
import pyopencl as cl
import numpy as np
from math import pow

clmem = cl.mem_flags
clqueue = cl.command_queue_properties



alphabet = np.array([ord(x) for x in ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h',
                                      'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p',
                                      'q', 'r', 's', 't', 'u', 'v', 'w', 'x',
                                      'y', 'z', '0', '1', '2', '3', '4', '5',
                                      '6', '7', '8', '9', 'A', 'B', 'C', 'D',
                                      'E', 'F', 'G', 'H', 'I', 'Y', 'K', 'L',
                                      'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T',
                                      'U', 'V', 'W', 'X', 'Y', 'Z', '_', '-']], dtype=np.uint8)

initStateBytes = np.array([0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476], dtype=np.uint32)


class MD5HashReverse:

    def __init__(self, profiling=False, useCPU=False):
        '''Initialize a new instance

        Initialize a new MD5HashReverse instance. The method is responsible of
        the initialization of the OpenCL environment. All additional structures
        needed for the attack (platform and device retrieving, OpenCL context
        and queues creation etc.) are initialized in this method.

        :param profiling: additional information, like attack time, throughput etc., will be printed

        :param useCPU: force the use of a CPU device instead of a GPU

        :throws: cl.Error exception if initialization fails

        '''
        self.profiling = profiling
        self.useCPU=useCPU
        self._initCL()

    
    def _initCL(self):
        platforms = cl.get_platforms();
        if len(platforms) == 0:
            raise cl.Error("Failed to find any OpenCL platform")
        
        # Get avaible GPU/CPU devices and choose the first one
        if self.useCPU:
            devices = platforms[0].get_devices(cl.device_type.CPU)
        else:
            devices = platforms[0].get_devices(cl.device_type.GPU)

        if len(devices) == 0:
            raise cl.Error("Could not find GPU devices")

        self.device = devices[0]

        # Create a context using the first GPU device found
        self.context = cl.Context([self.device])

        # Finally create the command queue
        if self.profiling:
            self.cmdQueue = cl.CommandQueue(self.context, self.device, 
                                            properties=clqueue.PROFILING_ENABLE)
        else:
            self.cmdQueue = cl.CommandQueue(self.context, self.device)


        kernelFile = open('md5.cl', 'r')
        kernelSrc = kernelFile.read()
        
        self.program = cl.Program(self.context, kernelSrc)
        self.program.build(devices=[self.device])



    def reverse(self, toReverseHash):
        '''Execute a brute-force attack against the given hash.

        Execute a brute-force attack against the given hash, trying to recover
        the original passphrase. All possible candidates with length between
        1 and 10 are tested. The original passphrase must be expressed using
        the English alphabet letters (both uppercase and lowercase), decimal
        digits and characters minus and underscore.
        
        :param toReverseHash: string of hexadecimal value representing the hash t attack

        :returns: recovered passphrase in case of success, and empty array otherwise

        '''
	# Subtract initial final values to input hash: it results in 4 less
        # addictions performed per thread
        toReverseHash.dtype = np.uint32
        toReverseHash -= initStateBytes
        toReverseHash.dtype = np.uint8

        alphBuff = cl.Buffer(self.context, clmem.READ_ONLY|clmem.COPY_HOST_PTR,
                             hostbuf=alphabet)
        
        inHashBuff = cl.Buffer(self.context, clmem.READ_ONLY|clmem.COPY_HOST_PTR,
                               hostbuf=toReverseHash)
        
        resBuff = cl.Buffer(self.context, clmem.WRITE_ONLY|clmem.COPY_HOST_PTR,
                            12, hostbuf=np.zeros(12, dtype=np.uint8))
        result = np.empty(12, dtype=np.uint8)

        # TODO: automatically select subspace size, maybe depending on
        # specified frequency of output info
        passOffset = 0
        passSpaceSize = int(pow(2,12))
        passSubSpaceSize = int(pow(2,25))

        
        for passLen in range(11)[1:10]:

            if self.profiling:
                print "\nTrying passphrases of length " + str(passLen)
                print "Passphrases space: {0} items".format(passSpaceSize)


            # Distinguish 3 different cases when launching the kernel for
            # breaking an hash of a given length:
            # - 1 the passphrases space is smaller than the workgroup size
            #   (it happens for passphrases of length 1, since we have
            #    2^6=64 possible passphrases, and the workgroup size is 256)
            if passSpaceSize < 128:
                #ev = self.program.md5hash(self.cmdQueue, (128,), (128,),
		ev = self.program.md5hash(self.cmdQueue, (256,), (256,),
                                          inHashBuff, np.uint64(passOffset),
                                          alphBuff, np.uint32(passLen),
                                          resBuff)
                if self.profiling:
                    pass
                
            # - 2 the passphrases space is smaller than a given passSubSpaceSize.
            #   The kernel is lauched once.
            elif passSpaceSize < passSubSpaceSize:
                #ev = self.program.md5hash(self.cmdQueue, (passSpaceSize,), (256,),
                ev = self.program.md5hash(self.cmdQueue, (passSpaceSize,),
                                          (128,),
                                          inHashBuff, np.uint64(passOffset),
                                          alphBuff, np.uint32(passLen),
                                          resBuff)
                if self.profiling:
                    pass

                # Check if we've found the hash
                cl.enqueue_copy(self.cmdQueue, result, resBuff)
                if result.any():
                    break

            # - 3 the passphrases space is greater than a given passSubSpaceSize.
            # In this case, the attach is going to take a lot of time. For the
            # reason, the whole passphrases space is broken into smaller chunks
            # of size passSubSpaceSize, and the kernel is executed once for
            # each chunk. This allows to print debug information meanwhile
            # an ongoing attack (i.e., after each chunk got attacked).
            # TODO: make the subspace size dependent on computer performance,
            # e.g. choose a chunk size such that each chunk gets attacked in
            # n seconds.
            else:
                for i in range(passSpaceSize/passSubSpaceSize):
                    #ev = self.program.md5hash(self.cmdQueue, (passSubSpaceSize,), (256,),
                    ev = self.program.md5hash(self.cmdQueue, (passSubSpaceSize,),
                                              (128,), inHashBuff,
                                              np.uint64(passOffset+(passSubSpaceSize*i)),
                                              alphBuff, np.uint32(passLen),
                                              resBuff)
                    cl.wait_for_events([ev,])
                    if self.profiling:
                        self._profile(ev, passSpaceSize, passSubSpaceSize, i+1)
                    

                    # Check if we've found the hash preimage
                    cl.enqueue_copy(self.cmdQueue, result, resBuff)
                    if result.any():
                        # Found!!!
                        break

                if result.any():
                    # If we're here, the hash preimage has already been found
                    break

            passOffset = 0
            passSpaceSize = int(pow(2, 6*(passLen+1)))

        return result


    def _profile(self, ev, passSpaceSize, passSubSpaceSize=0, subSpace=0):
        if passSubSpaceSize:
            percentage = (float(passSubSpaceSize)*subSpace)/passSpaceSize
            # TODO: use ev.profile.start instead of ev.profile.queued?
            execTime = (ev.profile.end - ev.profile.queued)*1e-9
            expectedExTime = execTime*((passSpaceSize/passSubSpaceSize)-subSpace)
            print "{0:.2f}% of passphrases space done in {1:.4f} seconds".format(percentage*100.0,
                                                                                 execTime)
            print "Expected remaining time: ~{0:.2f} seconds".format(expectedExTime)
        else:
            pass



def main():
    parser = argparse.ArgumentParser(description="MD5 hash reversing using OpenCL")
    parser.add_argument('--profiling', action='store_true')
    parser.add_argument('--use-cpu', action='store_true')
    parser.add_argument('hash')

    args = parser.parse_args()

    # TODO: perform additional check on input hash
    toRevHash = np.array([int(args.hash[x:x+2],16) for x in range(len(args.hash))[::2]],
                         dtype=np.uint8)

    reverser = MD5HashReverse(args.profiling, args.use_cpu)
    result = reverser.reverse(toRevHash)
    
    # TODO: nicer printing of result
    if result.any():
        passphrase = ''.join([chr(x) for x in result])
        print "Hash Reversed: [{0}]".format(passphrase[:passphrase.find('\x80')])
    else:
        print ":("



if __name__=='__main__':
    exit(main())
